#include <iostream>
#include <FansionML/KNNSolver>
#include <FansionML/Model>
#include <filesystem>
#include <fmt/core.h>
#include <thread>
#include "Model.h"

int getChoice(int mx) {
    fmt::print("选择：");
    int i;
    std::cin >> i;
    if (i <= 0 || i > mx) {
        fmt::print("不支持的选项：{}\n", i);
        exit(-1);
    }
    return i;
}

template<typename T>
T get(const std::string &desc) {
    std::cout << desc;
    T t;
    std::cin >> t;
    return t;
}

namespace fs = std::filesystem;

int main() {
    FansionML::Model *model = new Model();
    FansionML::Solver *solver;

    fmt::print("欢迎使用文本分类程序！\n");
    fmt::print("请选择读取数据的模式：\n1：加载并预处理数据；2：读取已预处理的数据\n");
    int i = getChoice(2);
    if (i == 1) {
        auto data_path = get<std::string>("请输入数据文件夹路径：");
        std::thread th{[&model, &data_path]() {
            model->preprocess(data_path.c_str());
        }};
        auto output_path = get<std::string>("请输入存储预处理结果文件夹的路径：");
        fmt::print("预处理中...\n");
        th.join();
        fmt::print("已完成预处理\n");
        model->savePreprocessData((fs::path(output_path) / "preprocess_data").c_str());
        fmt::print("已存储预处理结果\n");
    } else if (i == 2) {
        auto data_path = get<std::string>("请输入包含预处理结果的文件夹路径：");
        model->loadPreprocessData((fs::path(data_path) / "preprocess_data").c_str());
        fmt::print("已读取预处理结果：\n");
        model->explainData();
    }

    fmt::print("请选择分类器：\n1：KNN；2：Naïve Bayes\n");
    i = getChoice(2);
    if (i == 1) {
        int k = get<int>("请输入邻居数(k)：");
        solver = new FansionML::KNNSolver(k);
    } else if (i == 2) {
        fmt::print("暂未支持Naïve Bayes算法\n");
        exit(-1);
    }

    fmt::print("请选择计算方式：\n1：完整计算；2：分布式计算；3：合并分布式结果\n");
    i = getChoice(3);

    fmt::print("请输入存储训练结果文件夹的路径：");
    std::string output_path;
    std::cin >> output_path;
    model->setTestDataDir(output_path.c_str());

    if (i != 3) {
        int n = get<int>("请输入交叉验证的n值：");
        if (i == 1) {
            model->divideTrainTestData(n);
        } else if (i == 2) {
            fmt::print("请输入分布式分配值：\n");
            int n1 = get<int>(fmt::format("序号（1-{}）：", n)) - 1;
            int n2 = get<int>("起始进度（1-100）：") - 1;
            int n3 = get<int>("终止进度（1-100）：");
            model->divideTrainTestDataPart(n, n1, n2, n3);
        }

        solver->setThreadNum(get<int>("请输入并行计算的线程数（建议设为本机CPU所支持最大线程数减2）："));

        fmt::print("训练中...\n");
        model->trainAndEvaluate(solver);
        fmt::print("已完成训练\n已保存训练数据\n");
    } else {
        fmt::print("读取测试结果中...\n");
        model->mergeTestResult();
    }
}
